#include "audio_sr.h"
#include "esp_afe_sr_models.h"
#include <esp_err.h>
#include <freertos/FreeRTOS.h>
#include <freertos/task.h>
#include "bsp_codec.h"
// #include "object.h"
#include <string.h>
#include <stdlib.h>
#include <esp_log.h>

#define TAG "[AUDIO_PROCESSOR] SR"

struct audio_sr
{
    esp_afe_sr_iface_t *afe_handle;
    esp_afe_sr_data_t *afe_data;
    bool is_waken;
    vad_state_t last_state;

    void (*wakenet_callback)(void *);
    void *wakenet_callback_arg;

    void (*vad_callback)(void *, vad_state_t);
    void *vad_callback_arg;

    RingbufHandle_t output_buffer;
};

static void audio_sr_feed_task(void *arg)
{
    ESP_LOGI(TAG, "Feed task created");
    audio_sr_t *audio_sr = (audio_sr_t *)arg;
    esp_afe_sr_iface_t *afe_handle = audio_sr->afe_handle;
    esp_afe_sr_data_t *afe_data = audio_sr->afe_data;

    int feed_chunksize = afe_handle->get_feed_chunksize(afe_data);
    int feed_nch = afe_handle->get_feed_channel_num(afe_data);
    size_t feed_chunksize_bytes = feed_chunksize * feed_nch * sizeof(int16_t);
    int16_t *feed_buff = (int16_t *)malloc(feed_chunksize_bytes);

    // 用数据填满缓冲区
    while (1)
    {
        bsp_codec_read(feed_buff, feed_chunksize_bytes);
        afe_handle->feed(afe_data, feed_buff);
    }

    vTaskDelete(NULL);
}

static void audio_sr_fetch_task(void *arg)
{
    ESP_LOGI(TAG, "Fetch task created");
    audio_sr_t *audio_sr = (audio_sr_t *)arg;
    esp_afe_sr_iface_t *afe_handle = audio_sr->afe_handle;
    esp_afe_sr_data_t *afe_data = audio_sr->afe_data;

    int fetch_chunksize = afe_handle->get_fetch_chunksize(afe_data);
    int fetch_nch = afe_handle->get_fetch_channel_num(afe_data);
    size_t fetch_chunksize_bytes = fetch_chunksize * fetch_nch * sizeof(int16_t);

    while (1)
    {
        afe_fetch_result_t *result = afe_handle->fetch(afe_data);
        int16_t *raw_data = result->raw_data;
        vad_state_t vad_state = result->vad_state;
        wakenet_state_t wakeup_state = result->wakeup_state;

        if (wakeup_state == WAKENET_DETECTED)
        {
            // 检测到唤醒词
            audio_sr->is_waken = true;
            if (audio_sr->wakenet_callback)
            {
                audio_sr->wakenet_callback(audio_sr->wakenet_callback_arg);
            }
        }

        if (audio_sr->is_waken)
        {
            if (vad_state != audio_sr->last_state)
            {
                audio_sr->last_state = vad_state;
                if (audio_sr->vad_callback)
                {
                    audio_sr->vad_callback(audio_sr->vad_callback_arg, audio_sr->last_state);
                }
                if (vad_state == VAD_SPEECH && result->vad_cache_size > 0)
                {
                    // 将vad cache写入缓存
                    xRingbufferSend(audio_sr->output_buffer, result->vad_cache, result->vad_cache_size, 0);
                }
            }
        }
        else
        {
            if (audio_sr->last_state == VAD_SPEECH)
            {
                audio_sr->last_state = VAD_SILENCE;
                if (audio_sr->vad_callback)
                {
                    audio_sr->vad_callback(audio_sr->vad_callback_arg, audio_sr->last_state);
                }
            }
        }

        if (audio_sr->last_state == VAD_SPEECH)
        {
            // 将人声放到环形缓存
            xRingbufferSend(audio_sr->output_buffer, raw_data, fetch_chunksize_bytes, 0);
        }
    }

    vTaskDelete(NULL);
}

audio_sr_t *audio_sr_create(void)
{
    audio_sr_t *audio_sr = malloc(sizeof(audio_sr_t));
    assert(audio_sr);
    memset(audio_sr, 0, sizeof(audio_sr_t));

    srmodel_list_t *models = esp_srmodel_init("model");
    afe_config_t *afe_config = afe_config_init("M", models, AFE_TYPE_SR, AFE_MODE_LOW_COST);

    // 自定义afe配置
    afe_config->aec_init = false;
    afe_config->se_init = false;
    afe_config->ns_init = false;
    afe_config->vad_mode = VAD_MODE_3;
    // afe_config->vad_min_speech_ms = 64;
    // afe_config->vad_delay_ms = 64;
    afe_config->vad_min_noise_ms = 500;
    afe_config->wakenet_mode = DET_MODE_90;
    afe_config->memory_alloc_mode = AFE_MEMORY_ALLOC_MORE_PSRAM;

    // 获取句柄
    audio_sr->afe_handle = esp_afe_handle_from_config(afe_config);
    assert(audio_sr->afe_handle);
    // 创建实例
    audio_sr->afe_data = audio_sr->afe_handle->create_from_config(afe_config);
    assert(audio_sr->afe_data);

    return audio_sr;
}

void audio_sr_start(audio_sr_t *audio_sr)
{
    xTaskCreate(audio_sr_fetch_task, "fetch_task", 4096, audio_sr, 6, NULL);
    xTaskCreate(audio_sr_feed_task, "feek_task", 4096, audio_sr, 5, NULL);
}

void audio_sr_set_wakenet_callback(audio_sr_t *audio_sr, void (*callback)(void *), void *arg)
{
    audio_sr->wakenet_callback = callback;
    audio_sr->wakenet_callback_arg = arg;
}

void audio_sr_set_vad_callback(audio_sr_t *audio_sr, void (*callback)(void *, vad_state_t vad_state), void *arg)
{
    audio_sr->vad_callback = callback;
    audio_sr->vad_callback_arg = arg;
}

void audio_sr_set_output_buffer(audio_sr_t *audio_sr, RingbufHandle_t buf)
{
    audio_sr->output_buffer = buf;
}

void audio_sr_reset_wakenet(audio_sr_t *audio_sr)
{
    audio_sr->is_waken = false;
}
